package com.wuwangfu.state;

import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.functions.RichMapFunction;
import org.apache.flink.api.common.restartstrategy.RestartStrategies;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.KeyedStream;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;

import java.io.*;
import java.util.HashMap;

/**
 * @Author jcshen
 * @Date 2023-03-01
 * @PackageName:com.wuwangfu.state
 * @ClassName: KeyedStateByMap
 * @Description:
 * @Version 1.0.0
 *
 * 每个keyBy之后的subtask中定义一个hash map保存中间结果，定期将hash map中的数据持久化到磁盘，
 * 并且subtask出现异常重启，在open方法中可以读取磁盘中的文件，恢复历史数据。
 * （区别：每个定时器都是在subtask中定期执行，默认是JobManager定期向TaskManager发送checkpoint的rpc消息）
 */
public class KeyedStateByMap {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(2);
        //重启策略
        env.setRestartStrategy(RestartStrategies.fixedDelayRestart(3,5000));

        DataStreamSource<String> line = env.socketTextStream("localhost", 8888);
        SingleOutputStreamOperator<Tuple2<String, Integer>> maped = line.map(new MapFunction<String, Tuple2<String, Integer>>() {
            @Override
            public Tuple2<String, Integer> map(String value) throws Exception {
                String[] fields = value.split(",");
                if ("error".equals(fields[0])) {
                    throw new RuntimeException("出错啦");
                }
                return Tuple2.of(fields[0], Integer.parseInt(fields[1]));
            }
        });
        KeyedStream<Tuple2<String, Integer>, String> keyed = maped.keyBy(t -> t.f0);
        //
        keyed.map(new RichMapFunction<Tuple2<String, Integer>, Tuple2<String, Integer>>() {
            //存储中间结果的一个集合
            private HashMap<String,Integer> counter;

            @Override
            public void open(Configuration parameters) throws Exception {
                //初始化hash map 或恢复历史数据
                //获取当前subtask的编号
                int index = getRuntimeContext().getIndexOfThisSubtask();

                File ckfile = new File("D:/bigdata-ws/flink-parent/flink-14/ck" + index);
                if (ckfile.exists()){
                    FileInputStream fileStream = new FileInputStream(ckfile);
                    ObjectInputStream inputStream = new ObjectInputStream(fileStream);
                    counter = (HashMap<String, Integer>) inputStream.readObject();
                }else{
                    counter = new HashMap<>();
                }
                //简化：直接在当前subtask中启动一个定时器
                new Thread(new Runnable() {
                    @Override
                    public void run() {
                        while (true){
                            try {
                                Thread.sleep(10000);
                                //定期将hash map 写入文件
                                if (!ckfile.exists()){
                                    ckfile.createNewFile();
                                }
                                //将hash map 中的数据持久化到文件中
                                ObjectOutputStream outputStream = new ObjectOutputStream(new FileOutputStream(ckfile));
                                outputStream.writeObject(counter);
                                outputStream.flush();
                                outputStream.close();
                            } catch (Exception e) {
                                e.printStackTrace();
                            }
                        }
                    }
                }).start();
            }

            @Override
            public Tuple2<String, Integer> map(Tuple2<String, Integer> input) throws Exception {
                String word = input.f0;
                Integer  count= input.f1;

                Integer histCount = counter.get(word);
                if (histCount == null){
                    histCount = 0;
                }
                int sum = count + histCount;
                counter.put(word,sum);
                return Tuple2.of(word,sum);
            }
        }).print();

        env.execute();
    }
}
